# 导入函数库
from jqdata import *
import math
import scipy.optimize as sco
import numpy as np
import pandas as pd
import math

from matplotlib import pyplot as plt
from scipy.optimize import fmin, fminbound
from scipy import optimize
import datetime

EPS = 1e-15


# 初始化函数，设定基准等等
def initialize(context):
    # 设定基准
    set_benchmark('000300.XSHG')
    g.flag = 0
    g.security = '000854.OF'

    nd = context.current_dt.date()
    g.end_date = nd.strftime('%Y-%m-%d')
    start_year = str(int(nd.strftime('%Y')) - 1)
    g.start_date = start_year + nd.strftime('-%m-%d')

    # 开启动态复权模式(真实价格)
    set_option("use_real_price", True)

    ### 场外基金相关设定 ###
    # 设置账户类型: 场外基金账户
    set_subportfolios(
        [SubPortfolioConfig(context.portfolio.cash, 'open_fund')])
    # 设置赎回到账日
    set_redeem_latency(3000, 'open_fund')

    g.allstd = all_std()
    asp_arr = all_fund_select(g.allstd)
    g.stocks1000 = asp_arr[0:1000]
    asp_arr30 = asp_arr[0:20]
    g.stocks = asp_arr30

    # s_list      = get_meanVariance_list(asp_arr30)
    # print('s list',s_list)

    # fcorr_list  = get_fcorr_list(s_list)
    # print('f list',fcorr_list)
    # g.stocks    = list(set(s_list+fcorr_list ))

    # 运行函数（reference_security为运行时间的参考标的；传入的标的只做种类区分，因此传入'000300.XSHG'或'510300.XSHG'是一样的）
    run_daily(
        before_market_open,
        time='9:30',
        reference_security='000300.XSHG')
    # 开盘时运行
    run_daily(market_open, time='open', reference_security='000300.XSHG')
    # 收盘后运行
    run_daily(
        after_market_close,
        time='after_close',
        reference_security='000300.XSHG')


def all_std():
    all_fund = get_all_securities(types=['open_fund'])
    funds = list(all_fund.index)
    df = get_df2(funds)
    rate = (df - df.shift(1)) / df.shift(1)
    # print(df.columns.values)
    allstd = rate.std()
    # print(rate.std())
    # print(allstd['001984.OF'])
    return allstd


def get_fcorr_list(s_list):
    fcorr_list = []
    for item in s_list:
        code = f_corr(item)
        # print(item,code)
        if code != "":
            fcorr_list.append(code)
    return fcorr_list


def cal_portfolio(args):
    means, covs = args

    def obj(weights):
        weights = np.array(weights)
        port_mean = np.dot(weights.T, means)
        port_var = np.sqrt(np.dot(weights.T, np.dot(covs, weights)))
        return np.array([port_mean, port_var, port_mean / (EPS + port_var)])

    return obj


def cal_neg_sharpe(args):
    means, covs = args

    def obj(weights):
        weights = np.array(weights)
        port_mean = np.dot(weights.T, means)
        port_var = np.sqrt(np.dot(weights.T, np.dot(covs, weights)))
        return -np.array(port_mean / (EPS + port_var))

    return obj


def cal_mean(args):
    means, covs = args

    def obj(weights):
        weights = np.array(weights)
        port_mean = np.dot(weights.T, means)
        return np.array(port_mean)

    return obj


def cal_variance(args):
    means, covs = args

    def obj(weights):
        weights = np.array(weights)
        port_var = np.sqrt(np.dot(weights.T, np.dot(covs, weights)))
        return np.array(port_var)

    return obj


class meanVariance(object):
    def __init__(self, returns):
        self.means = np.array(returns.mean())
        self.covs = np.array(returns.cov())
        self.n_assets = len(self.means)

    def sharpe(self, weights):
        obj = cal_neg_sharpe((self.means, self.covs))
        return -obj(weights)

    def mean_return(self, weights):
        obj = cal_mean((self.means, self.covs))
        return obj(weights)

    def variance(self, weights):
        obj = cal_variance((self.means, self.covs))
        return obj(weights)

    def opt_min_sharpe(self, add_cons=None, add_bnds=None):
        if add_cons is not None:
            cons = tuple(
                [{'type': 'eq', 'fun': lambda x: np.sum(x) - 1}] + add_cons)
        else:
            cons = ({'type': 'eq', 'fun': lambda x: np.sum(x) - 1})

        if add_bnds is not None:
            bnds = [(0, 1) for x in range(self.n_assets)]
            for i, bnd in add_bnds:
                bnds[i] = bnd
            bnds = tuple(bnds)
        else:
            bnds = tuple((0, 1) for x in range(self.n_assets))

        opt_ret = sco.minimize(cal_neg_sharpe((self.means, self.covs)), np.ones(
            self.n_assets) / self.n_assets, method='SLSQP', bounds=bnds, constraints=cons)

        return opt_ret

    def opt_min_variance(self, add_cons=None, add_bnds=None):
        if add_cons is not None:
            cons = tuple(
                [{'type': 'eq', 'fun': lambda x: np.sum(x) - 1}] + add_cons)
        else:
            cons = ({'type': 'eq', 'fun': lambda x: np.sum(x) - 1})

        if add_bnds is not None:
            bnds = [(0, 1) for x in range(self.n_assets)]
            for i, bnd in add_bnds:
                bnds[i] = bnd
            bnds = tuple(bnds)
        else:
            bnds = tuple((0, 1) for x in range(self.n_assets))

        opt_ret = sco.minimize(cal_variance((self.means, self.covs)), np.ones(
            self.n_assets) / self.n_assets, method='SLSQP', bounds=bnds, constraints=cons)

        return opt_ret


def get_stocks_prices(stocks):
    # stocks = ['000854.OF', '001216.OF']
    try:
        prices = get_extras(
            'acc_net_value',
            stocks,
            start_date=g.start_date,
            end_date=g.end_date,
            df=True)
        # prices = get_extras('acc_net_value', stocks, count=g.count, df=True)
        return prices
    except Exception as r:
        print('未知错误 %s' % (r))
    return pd.DataFrame()


def get_meanVariance_list(stocks):
    # stocks = ['000854.OF', '001216.OF']
    prices = get_stocks_prices(stocks)
    returns = (prices - prices.shift(1)) / prices.shift(1)
    # log.info('sum')
    # log.info(returns.sum())
    # log.info('std')
    # log.info(returns.std())
    # log.info('mean')
    # log.info(returns.mean())
    # log.info('corr')
    # log.info(returns[stocks[0]].corr(returns[stocks[1]]))
    # log.info('cov')
    # log.info(returns[stocks[0]].cov(returns[stocks[1]]))

    mv = meanVariance(returns)  # 创建投资组合模型
    print('finished meanVariance')
    means = mv.means
    covs = mv.covs
    print('finished meanVariance2')
    fund0 = means[0]
    add_cons = [{"type": "ineq", "fun": lambda x: x[0] * fund0 - 0.2 * fund0}]
    # add_bnds = [(2, (0.1, 0.3))]
    # add_cons = None
    add_bnds = None
    ret = mv.opt_min_sharpe(add_cons=add_cons, add_bnds=add_bnds)
    print('finished meanVariance3')
    if g.flag == 1:
        return
    if ret['success']:
        weights = ret['x']
        # log.info('\n', ret)
        # log.info("==================")
        # log.info("sharpe rate", mv.sharpe(weights))
        # log.info("mean returns", mv.mean_return(weights))
        # log.info("weights", weights)
        # log.info("sum of weights", np.sum(list(weights)))
        # log.info("fund 0 return", weights[0] * means[0])
        # log.info("==================")

        # total_value = context.portfolio.total_value  # 获取总资产
        s_list = []
        for i, stock in enumerate(stocks):
            # value = total_value * weights[i] # 确定每个标的的权重
            # purchase(stock, value)
            if weights[i] > 0:
                s_list.append(stocks[i])
            # order_target_value(stock, value)  # 调整标的至目标权重
        # g.flag=1
        return s_list
    else:
        return []
        log.info("组合优化失败")


def test(context):
    stocks = g.stocks
    prices = get_stocks_prices(stocks)
    returns = (prices - prices.shift(1)) / prices.shift(1)

    log.info('sum')
    log.info(returns.sum())
    log.info('std')
    log.info(returns.std())
    log.info('mean')
    log.info(returns.mean())
    log.info('corr')
    log.info(returns[stocks[0]].corr(returns[stocks[1]]))
    log.info('cov')
    log.info(returns[stocks[0]].cov(returns[stocks[1]]))

    mv = meanVariance(returns)  # 创建投资组合模型
    means = mv.means
    covs = mv.covs
    fund0 = means[0]
    n_assets = mv.n_assets
    var0 = mv.variance(np.ones(n_assets) / n_assets)
    add_cons = [{"type": "ineq", "fun": lambda x: x[0] * fund0 - 0.5 * fund0}]
    # add_cons = [{"type": "ineq", "fun": lambda x: -mv.variance(x)+2.0*var0}]
    # add_bnds = [(0, (0,0.6)), (2, (0.1, 1))]
    # add_cons = None
    add_bnds = None
    ret = mv.opt_min_sharpe(add_cons=add_cons, add_bnds=add_bnds)
    weights = ret['x']
    log.info('\n', ret)
    log.info("==================")
    log.info("sharpe rate", mv.sharpe(weights))
    log.info("mean returns", mv.mean_return(weights))
    log.info("weights", weights)
    log.info("sum of weights", np.sum(list(weights)))
    log.info("fund 0 return", weights[0] * means[0])
    log.info("==================")


# 开盘前运行函数
def before_market_open(context):
    print('调仓日期：%s' % context.current_dt.date())

    g.buy_list = g.stocks


# 开盘时运行函数
def market_open(context):
    if g.flag == 1:
        return
    # 将不在股票池中的股票卖出
    sell_list = set(context.portfolio.positions.keys()) - set(g.buy_list)
    for stock in sell_list:
        order_target_value(stock, 0)

    stocks = g.buy_list
    log.info("stock list", stocks)
    prices = get_stocks_prices(stocks)
    returns = (prices - prices.shift(1)) / prices.shift(1)

    mv = meanVariance(returns)  # 创建投资组合模型
    means = mv.means
    covs = mv.covs

    fund0 = means[0]
    add_cons = [{"type": "ineq", "fun": lambda x: x[0] * fund0 - 0.2 * fund0}]
    # add_bnds = [(2, (0.1, 0.3))]
    # add_cons = None
    add_bnds = None
    ret = mv.opt_min_sharpe(add_cons=add_cons, add_bnds=add_bnds)

    if ret['success']:
        weights = ret['x']
        # log.info('\n', ret)
        # log.info("==================")
        # log.info("sharpe rate", mv.sharpe(weights))
        # log.info("mean returns", mv.mean_return(weights))
        log.info("weights", weights)
        log.info("sum of weights", np.sum(list(weights)))
        # log.info("fund 0 return", weights[0] * means[0])
        # log.info("==================")

        total_value = context.portfolio.total_value  # 获取总资产
        for i, stock in enumerate(stocks):
            value = total_value * weights[i]  # 确定每个标的的权重
            if value > 0:
                purchase(stock, value)
            # order_target_value(stock, value)  # 调整标的至目标权重
        g.flag = 1
    else:
        log.info("组合优化失败")


# 根据年化排序
def all_fund_select(allstd):
    all_fund = get_all_securities(types=['open_fund'])
    funds = list(all_fund.index)
    all_returns = {}
    all_returns_list = []
    i = 0
    # funds= ['001832.OF', '001984.OF']
    for item in funds:
        i = i + 1
        # print('fund:'+str(i))
        try:
            res = get_return2(item)
            df = get_df(item)
            # df2=get_df2(item)
            # rate=(df2-df2.shift(1))/df2.shift(1)
            # std=rate[item].std()
            std = allstd[item]
            # print(item ,std)
            # print(df)
            if len(df) > 0:
                dd = MaxDrawdown(df)
                if res > -100:
                    # if item in ['000854.OF', '001216.OF', '000934.OF', '110022.OF', '519674.OF', '001984.OF']:
                    #     print('------------------')
                    #     print(item+':rt '+str(res))
                    #     print(item+':dd '+str(dd))
                    #     print(item+':std '+str(std))

                    if res > 0.08 and dd > 0.02:
                        all_returns[item] = res / (dd * 0.8 + 0.8 * std)
                        # all_returns[item]=res/std
                        # all_returns[item]=res/dd*10+res/(std*0.2)
                        # if item in ['000854.OF', '001216.OF', '000934.OF', '110022.OF', '519674.OF', '001984.OF']:
                        #     print(item,all_returns[item])
            else:
                print('df empty')
        except Exception as r:
            print('未知错误 %s' % (r))
            print(item)
            # print(df)
    print('finished for')
    aps = sorted(all_returns.items(), key=lambda d: d[1], reverse=True)
    # log.info(aps[0:20])
    # log.info(aps)
    # log.info('guava')

    apss = ""
    aps_arr = []
    for item in aps:
        apss = apss + "'" + str(item[0]) + "',"
        aps_arr.append(item[0])

    # print(aps_arr)
    return aps_arr


def get_return2(stocks):
    # stocks = ['000854.OF', '001216.OF']
    try:
        df = get_extras(
            'acc_net_value',
            stocks,
            start_date=g.start_date,
            end_date=g.end_date,
            df=False)
        # df = get_extras('acc_net_value',stocks, count= g.count, df=False)

        arr = df[stocks]
        # print(arr)
        return (arr[-1] - arr[0]) / arr[0]
    except Exception as r:
        # print(stocks)
        print('error %s' % (r))
        print(df)
        return -100


def get_df2(stocks):
    # stocks = ['000854.OF', '001216.OF']
    try:
        df = get_extras(
            'acc_net_value',
            stocks,
            start_date=g.start_date,
            end_date=g.end_date,
            df=True)
        # df = get_extras('acc_net_value',stocks, count= g.count, df=True)
        # rate=(df-df.shift(1))/df.shift(1)
        return df
    except Exception as r:
        print('未知错误 %s' % (r))
    return pd.DataFrame()


def get_df(stocks):
    # stocks = ['000854.OF', '001216.OF']
    try:
        df2 = get_extras(
            'acc_net_value',
            stocks,
            start_date=g.start_date,
            end_date=g.end_date,
            df=True)
        # df2 = get_extras('acc_net_value',stocks, count= g.count, df=True)
        df2 = df2.fillna(method='ffill')
        df = df2[stocks].tolist()
        # print(df)
        return df
    except Exception as r:
        print('未知错误 %s' % (r))
        return pd.DateFrame()


def max_sharp_g(stocks):
    # print('in max_sharp_g')
    g.stocks = stocks
    df = get_df2(g.stocks)
    rate = (df - df.shift(1)) / df.shift(1)
    # print(df.columns.values)
    if stocks not in df.columns.values:
        return 0
    # print(df['519062.OF'])
    l1 = rate[stocks[0]]
    l2 = rate[stocks[1]]
    std1 = g.allstd[stocks[0]]
    # print('std:'+str(std1))
    std2 = g.allstd[stocks[1]]
    the_return = rate.sum()
    # df=get_df(g.stocks)
    # dd = maxdrawdown(df)
    if False:
        log.info('sum')
        log.info(rate.sum())
        log.info('std')
        log.info(rate.std())
        log.info('mean')
        log.info(rate.mean())
        log.info('corr')
        log.info(rate[stocks[0]].corr(rate[stocks[1]]))
        log.info('cov')
        log.info(rate[stocks[0]].cov(rate[stocks[1]]))

    # fminbound
    x = np.linspace(0.1, 1, num=10)
    y = []
    for item in x:
        # fstd= f_sharp(item,l1,l2,l1.std(),l2.std(),l1.corr(l2),the_return)
        # print("============b")
        # print(item,std1,std2,l1.corr(l2),the_return)
        # print("============e")
        fstd = f_sharp(item, std1, std2, l1.corr(l2), the_return)
        # print(fstd)
        y.append(fstd)

    # print(y)
    # print(max(y))
    min_index = np.argmax(y)
    # print(g.stocks,x[min_index])
    return max(y), x[min_index]


def f_std(x, l1, l2):
    aa = pow(l1.std() * x, 2) + pow(l2.std() * (1 - x), 2)
    +2 * l1.corr(l2) * (1 - x) * x * l1.std() * l2.std()
    ac = math.sqrt(aa)
    return ac


def f_sharp(x, l1std, l2std, corr12, the_return):
    try:
        # print(corr12)
        # if corr12>-0.2:
        #     return 0
        aa = pow(l1std * x, 2) + pow(l2std * (1 - x), 2) + \
            2 * corr12 * (1 - x) * x * l1std * l2std
        ac = math.sqrt(aa)

        returns = the_return[g.stocks[0]] * \
            x + the_return[g.stocks[1]] * (1 - x)

        if returns < 0.3:
            return 0
        # print('new return:',returns)
        # print(ac)
        # print(returns/ac)
        # print(x,'-----')
        return returns / ac
    except Exception as r:
        # print('未知错误 %s' %(r))
        return 0
    # return 0


def f_corr(base_fund='001216.OF'):
    # all_fund= get_all_securities(types=['open_fund'])
    # stocks  = list(all_fund.index)
    stocks = g.stocks1000
    # print(base_fund+' :in fcorr')
    max_sharp = 0
    code = ""
    for item in stocks:
        stock = [base_fund, item]
        # log.info('f_corr item:',item)

        try:
            v_sharp, rate = max_sharp_g(stock)
        except Exception as r:
            v_sharp = 0

        if v_sharp > max_sharp:
            max_sharp = v_sharp
            code = item
            # print('sharp rate:',max_sharp)
            # print(stock)
            # print('-------------')
    return code


# def market_open(context):
#     all_fund=get_all_securities(types=['open_fund'])
#     #for index, rows in all_fund.iterrows():
#     #    if rows['name']=='海富通阿尔法A':
#     #        log.info(rows)
#     # log.info(all_fund['519062.OF'])
#     funds = list(all_fund.index)
#     # security_info= get_security_info('502050.XSHG')
#     # log.info(security_info.start_date)
#     res= get_return2(g.stocks[0])
#     log.info(res)
#     # log.info(funds)
#     # 设置场外基金标的为景顺沪深300('000311.OF')
#     s = g.stocks[0]

#     # 获取基金信息
#     fund_info = get_fund_info(s)
#     log.info("基金名称：",fund_info['fund_name'])

#     # 确定时间是周几
#     weekday = context.current_dt.isoweekday()
#     log.info("今天是周 %s" % weekday)
#     #得到当前以持有的股票池
#     present = context.portfolio.positions
#     if s in present:
#         log.info('in positions')
#     #得到当前股票池中的股票代码
#     stocksnow = present.keys()
#     log.info(len(stocksnow))
#     #如果当前股票池中有股票，得到每个股票对应的仓位市值，数据结构为字典
#     if g.flag== 0:
#         o = purchase(g.stocks[0], g.x*100000)
#         o = purchase(g.stocks[1], (1-g.x)*100000)
#         log.info(context.portfolio.subportfolios[0])
#         log.info(o)
#         g.flag=1
#     # 赎回基金
#     # elif weekday == 3:
#     #     o1 = redeem(s, 4000)
#     #     log.info(o1)
#     # elif weekday == 4:
#     #     o2 = redeem(s, 3000)
#     #     log.info(o2)

# ## 收盘后运行函数
def after_market_close(context):
    # 查看融资融券账户相关相关信息(更多请见API-对象-SubPortfolio)
    p = context.portfolio.subportfolios[0]
    log.info('- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -')
    log.info('查看场外基金账户相关相关信息(更多请见API-对象-SubPortfolio)：')
    log.info('场外基金持有份额：', p.long_positions['519062.OF'].closeable_amount)
    log.info('账户所属类型：', p.type)
    log.info('##############################################################')


def maxdrawdown(arr):
    i = np.argmax((np.maximum.accumulate(arr) - arr) /
                  np.maximum.accumulate(arr))  # end of the period
    j = np.argmax(arr[:i])  # start of period
    return (1 - arr[i] / arr[j])


def MaxDrawdown(return_list):
    # print('最大回撤率')
    i = np.argmax(
        (np.maximum.accumulate(return_list) -
         return_list) /
        np.maximum.accumulate(return_list))  # 结束位置
    # print('最大回撤率'+ str(i))
    if i == 0:
        return 0
    j = np.argmax(return_list[:i])  # 开始位置
    # print('最大回撤率'+ str(j))
    # print(return_list[j])
    # cha = (return_list[j] - return_list[i])
    # print(cha)
    # print( return_list[197])
    # print('len:',len( return_list))
    # print( return_list[-1])
    # print( return_list[0])
    # print( return_list)

    return (return_list[j] - return_list[i]) / (return_list[j])
